
#include "LinearBlock.hpp"
#include "ReLu.hpp"

namespace qing {
    class NeuralNetwork {
    private:
        std::vector<std::unique_ptr<Layer>> layers;

    public:
        // 添加层
        void add_layer(std::unique_ptr<Layer> layer) {
            // 检查层之间的维度是否匹配
            if (!layers.empty()) {
                int prev_output = layers.back()->get_output_dim();
                int current_input = layer->get_input_dim();
                if (prev_output != current_input) {
                    throw std::invalid_argument("Layer dimension mismatch: " + 
                        std::to_string(prev_output) + " -> " + std::to_string(current_input));
                }
            }
            layers.push_back(std::move(layer));
        }

        // 前向传播
        std::vector<double> forward(const std::vector<double>& input) {
            std::vector<double> output = input;
            for (auto& layer : layers) {
                output = layer->forward(output);
            }
            return output;
        }

        // 反向传播
        std::vector<double> backward(const std::vector<double>& grad_output) {
            std::vector<double> grad = grad_output;
            for (auto it = layers.rbegin(); it != layers.rend(); ++it) {
                grad = (*it)->backward(grad);
            }
            return grad;
        }

        // 更新所有参数
        void update_parameters(double learning_rate) {
            for (auto& layer : layers) {
                layer->update_parameters(learning_rate);
            }
        }

        // 设置训练模式
        void set_training(bool training) {
            for (auto& layer : layers) {
                layer->set_training(training);
            }
        }

        // 获取总参数数量
        int get_total_parameters() const {
            int total = 0;
            for (const auto& layer : layers) {
                total += layer->get_parameter_count();
            }
            return total;
        }

        // 打印网络结构
        void print_architecture() const {
            std::cout << "Neural Network Architecture:" << std::endl;
            std::cout << "Total parameters: " << get_total_parameters() << std::endl;
            std::cout << "Layers:" << std::endl;
            
            for (size_t i = 0; i < layers.size(); ++i) {
                std::cout << "  " << i + 1 << ". " << layers[i]->get_layer_type()
                        << ": " << layers[i]->get_input_dim() 
                        << " -> " << layers[i]->get_output_dim() 
                        << " (params: " << layers[i]->get_parameter_count() << ")" << std::endl;
            }
        }
    };
}